#define _CRT_SECURE_NO_WARNINGS

#include <stdio.h>
#include <io.h>

#include <Windows.h>

#include "utils.h"
#include "xml_parse.h"

char msg_buf[1024];

struct sector_conf{
	char *name;
	int sector;
	struct sector_conf *next;
};

void free_conf(struct sector_conf *head)
{
	struct sector_conf *node;
	
	node = head;
	while (node != NULL) {
		struct sector_conf *tmp = node->next;
		if (node->name)
			free(node->name);
		free(node);
		node = tmp;
	}
}

struct sector_conf *process_conf(char *file_name)
{
	HANDLE sector_file;
	char buffer[1024];
	char line[1024];
	DWORD num_read;
	struct sector_conf *head = NULL, *node = NULL;
	DWORD cur_pos = 0, line_pos = 0;
	DWORD i;

	sector_file = CreateFile(file_name, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
	if (sector_file == INVALID_HANDLE_VALUE) {
		printf("open file %s failed: %s\n", file_name, last_error(msg_buf, MSG_BUF_SIZE));
		return NULL;
	}
	head = (struct sector_conf *)malloc(sizeof(struct sector_conf));
	if (head == NULL) {
		errmsg("out of memory!\n");
		goto err_out;
	}
	head->name = NULL;
	head->next = NULL;

	node = head;

	while (ReadFile(sector_file, buffer, 10, &num_read, NULL) && num_read > 0) {
		cur_pos = 0;
		while (cur_pos < num_read) {
			if (buffer[cur_pos] == '\n') {
				line[line_pos] = '\0';
				//new line copied
				node->next = (struct sector_conf *)malloc(sizeof(struct sector_conf));
				if (node->next == NULL) {
					errmsg("out of memory!\n");
					goto err_out;
				}
				node = node->next;
				node->next = NULL;
				i = 0;
				while (i < line_pos && line[i] != 0) {
					if (line[i] == ':') {
						line[i] = '\0';
						node->name = _strdup(line);
						node->sector = atoi(line + i + 1);
						break;
					}
					i++;
				}
				line_pos = 0;
			} else if (buffer[cur_pos] != '\r') {
				line[line_pos++] = buffer[cur_pos];
			}
			cur_pos++;
		}
	}

	CloseHandle(sector_file);
	return head;

err_out:
	free_conf(head);

	CloseHandle(sector_file);
	return NULL;
}

struct sector_conf *find_part(struct sector_conf *all_conf, char *part_name)
{
	struct sector_conf *node, *head = NULL, *part = NULL, *tmp;

	head = (struct sector_conf *)malloc(sizeof(struct sector_conf));
	if (head == NULL) {
		errmsg("out of memory!\n");
		return NULL;
	}
	head->name = NULL;
	head->next = NULL;

	part = head;
	node = all_conf->next;
	while (node) {
		if (strstr(node->name, part_name)) {
			part = (struct sector_conf *)malloc(sizeof(struct sector_conf));
			if (part == NULL) {
				errmsg("out of memory!\n");
				goto err_out;
			}
			part->name = _strdup(node->name);
			part->sector = node->sector;
			tmp = head;
			while (tmp->next && (tmp->next->sector < node->sector))
				tmp = tmp->next;
			part->next = tmp->next;
			tmp->next = part;
		}
		node = node->next;
	}

	return head;

err_out:
	free_conf(head);

	return NULL;
}

int write_file(struct sector_conf *head, char* in_folder, char *out_file_name)
{
	struct sector_conf *node;
	int base_sector;
	char in_file[MAX_PATH];
	HANDLE h_out_file, h_in_file;
	char buffer[1024];
	DWORD num_read, num_written;

	node = head->next;
	if (node == NULL)
		return -1;

	base_sector = node->sector;

	h_out_file = CreateFile(out_file_name, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
	if (h_out_file == INVALID_HANDLE_VALUE) {
		printf("create file %s failed: %s\n", out_file_name, last_error(msg_buf, MSG_BUF_SIZE));
		return -1;
	}
	while(node) {
		sprintf(in_file, "%s\\APPS\\%s", in_folder, node->name);
		h_in_file = CreateFile(in_file, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
		if (h_in_file == INVALID_HANDLE_VALUE) {
			printf("open file %s failed: %s\n", out_file_name, last_error(msg_buf, MSG_BUF_SIZE));
			CloseHandle(h_out_file);
			return -1;
		}
		__int64 offset = (__int64)(node->sector - base_sector) * SECTOR_SIZE;
		
		if (seek_file(h_out_file, offset, FILE_BEGIN) < 0){
			CloseHandle(h_out_file);
			return -1;
		}

		while (ReadFile(h_in_file, buffer, 1024, &num_read, NULL) && num_read > 0) {
			WriteFile(h_out_file, buffer, num_read, &num_written, NULL);
			if (num_written != num_read) {
				errmsg("WriteFile error: %s\n", last_error(msg_buf, MSG_BUF_SIZE));
				CloseHandle(h_out_file);
				return -1;
			}
		}
		node = node->next;
		CloseHandle(h_in_file);
	}
	
	CloseHandle(h_out_file);

	return 0;
}

int merge_img_conf(char *folder)
{
	char file_name[MAX_PATH];
	struct sector_conf *conf;
	struct sector_conf *system_head = NULL, *udata_head = NULL;

	if (_access(folder, 0) != 0) {
		errmsg("folder %s does not exist.\n", folder);
		return -1;
	}

	sprintf(file_name, "%s\\%s", folder, "SectorPosition");
	conf = process_conf(file_name);
	if (conf == NULL)
		return -1;

	system_head = find_part(conf, "system");
	if (system_head == NULL)
		goto err_out1;

	udata_head = find_part(conf, "udata");
	if (udata_head == NULL)
		goto err_out2;
	
	write_file(system_head, folder, "system.img");
	write_file(udata_head, folder, "udata.img");

	free_conf(udata_head);
err_out2:
	free_conf(system_head);
err_out1:
	free_conf(conf);

	return 0;
}

int merge_img_xml(NL& node_list, const char *folder, const char *image_name) {
	
	char file_name[MAX_PATH];
	// system_nodes.sort(); // we need a sort to sort the start address
	string val;
	HANDLE h_out_file, h_in_file;
	DWORD sector, base_sector, num_read, num_written;
	__int64 offset;
	char buffer[1024];
	
	if (node_list.size() <= 0)
		return -1;
	
	sprintf(file_name, "%s\\out\\%s.img", folder, image_name);
	h_out_file = CreateFile(file_name, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
	if (h_out_file == INVALID_HANDLE_VALUE) {
		errmsg("create file %s failed: %s\n", file_name, last_error(msg_buf, MSG_BUF_SIZE));
		return -1;
	}
	NLI i = node_list.begin();
	(*i)->get_val("start_sector", val);
	if (val.empty()) {
		errmsg("can't get start_sector\n");
		return -1;
	}
	base_sector = atoi(val.c_str());
	for (NLI i = node_list.begin(); i != node_list.end(); i++) {
		(*i)->get_val("filename", val);
		if (!val.empty()) {
			sprintf(file_name, "%s\\%s", folder, val.c_str());
			h_in_file = CreateFile(file_name, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
			if (h_in_file == INVALID_HANDLE_VALUE) {
				printf("open file %s failed: %s\n", file_name, last_error(msg_buf, MSG_BUF_SIZE));
				CloseHandle(h_out_file);
				return -1;
			}
			(*i)->get_val("start_sector", val);
			if (val.empty()) {
				errmsg("can't get start_sector\n");
				CloseHandle(h_in_file);
				CloseHandle(h_out_file);
				return -1;
			}
			sector = atoi(val.c_str());
			offset = (__int64)(sector - base_sector) * SECTOR_SIZE;
			if (seek_file(h_out_file, offset, FILE_BEGIN) < 0){
				CloseHandle(h_in_file);
				CloseHandle(h_out_file);
				return -1;
			}
			while (ReadFile(h_in_file, buffer, 1024, &num_read, NULL) && num_read > 0) {
				WriteFile(h_out_file, buffer, num_read, &num_written, NULL);
				if (num_written != num_read) {
					errmsg("WriteFile error: %s\n", last_error(msg_buf, MSG_BUF_SIZE));
					CloseHandle(h_in_file);
					CloseHandle(h_out_file);
					return -1;
				}
			}
			CloseHandle(h_in_file);
		} else {
			(*i)->get_val("num_partition_sectors", val);
			if (!val.empty()) {
				sector = atoi(val.c_str());
				offset = ((__int64)sector * SECTOR_SIZE);
				if (seek_file(h_out_file, offset, FILE_BEGIN) < 0){
					CloseHandle(h_out_file);
					return -1;
				}
				SetEndOfFile(h_out_file);
			}
		}
	}
	return 0;
}
;

#if 0
BEAB 7D0
C67A 3E8
CA61 FA0
DA00 106A
1440F 1800
15C0E 1800
1740D 1800
18C0C FA000
112CB0 248B00
35B70A 5000
360709 32000
392708 6A10
399117 C000
3A5116 37A5A8
71F6BD 7C800
71FE84 0

#endif

void usage()
{
	printf("mergeImg.exe x D:\\msm\n");
	printf("mergeImg.exe c D:\\msm\n");
	return;
}
int main(int argc, char* argv[])
{
	string op, folder;

//#define DEBUG_ARGS
#ifdef DEBUG_ARGS
	op = "c";
	folder = "D:\\images\\20141216";
#else
	if (argc != 3) {
		usage();
		return -1;
	}

	op = argv[1];
	if (op != "x" && op != "c") {
		usage();
		return -1;
	}
	folder = argv[2];
#endif

	if (_access(folder.c_str(), 0) != 0) {
		errmsg("folder %s does not exist.\n", folder.c_str());
		return -1;
	}
	
	string xml_file_name = "rawprogram0.xml";
	string xml_file = folder + "\\" + xml_file_name;
	if (_access(xml_file.c_str(), 0) != 0) {
		xml_file_name = "rawprogram_unsparse.xml";
		xml_file = folder + "\\" + xml_file_name;
		if (_access(xml_file.c_str(), 0) != 0) {
			errmsg("can't find \"rawprogram0.xml\" or \"rawprogram_unsparse.xml\"");
		}
	}

	xml_parser xp(xml_file);

	xml_node *root = xp.get_root_node();
	if (root == NULL || root->get_first_child() == NULL)
		return -1;

	NL data_nodes, program_nodes;
	NL system_nodes, user_nodes;

	xp.find_node_byname(root, "data", data_nodes);

	for (NLI i = data_nodes.begin(); i != data_nodes.end(); i++) {
		if ((*i)->get_type() == END_NODE) continue;
		xp.find_node_byname(*i, "program", program_nodes);
	};

	for (NLI i = program_nodes.begin(); i != program_nodes.end(); i++) {
		if ((*i)->get_type() == END_NODE) continue;
		string val;
		(*i)->get_val("label", val);
		transform(val.begin(), val.end(), val.begin(), toupper);
		if (val == "SYSTEM") {
			system_nodes.push_back(*i);
		} else if (val == "USERDATA") {
			user_nodes.push_back(*i);
		}
	};
		
	if (op == "x") {
		string out_folder = folder + "\\out";
		if (_access(out_folder.c_str(), 0) != 0) {
			if (!CreateDirectory(out_folder.c_str(), NULL)) {
				errmsg("create directory %s failed\n", out_folder.c_str());
				return -1;
			}
		}

		merge_img_xml(system_nodes, folder.c_str(), "system");
		merge_img_xml(user_nodes, folder.c_str(), "userdata");
	} else {
		string image_file;

		image_file = folder + "\\out\\system.img";
		if (split_file(image_file.c_str(), system_nodes) != 0) {
			errmsg("split system.img failed\n");
			return -1;
		}

		image_file = folder + "\\out\\userdata.img";
		if (split_file(image_file.c_str(), user_nodes) != 0) {
			errmsg("split userdata.img failed\n");
			return -1;
		}

		if (xp.write_xml(folder + "\\out\\" + xml_file_name) != 0) {
			errmsg("write xml file failed\n");
			return -1;
		}
		
		//infomsg("delete old unsparse files\n");
#if 1
		string cmd;
		cmd = "del " + folder + "\\system* 2> nul";
		system(cmd.c_str());
		cmd = "del " + folder + "\\userdata* 2> nul";
		system(cmd.c_str());
		cmd = "move " + folder + "\\out\\*_* " + folder + " >nul";
		system(cmd.c_str());
		cmd = "move " + folder + "\\out\\" + xml_file_name + " " + folder + " > nul";
		system(cmd.c_str());
		//cmd = "del " + folder + "\\out\\*.img 2 > nul";
		//system(cmd.c_str());
		//cmd = "rmdir " + folder + "\\out 2 > nul";
		//system(cmd.c_str());
#endif
	}
	return 0;
}
